#!/usr/bin/env python
'''
This module provides function to extract 
sequence, tag (Name) and Names dictionary
print sequence in the fasta file 

Created on Sep 17, 2009
@author: yanghoch
'''
import string

def reverseCompliment( Seq ):
    def complement_seq( seq ):
        complement = {'A': 'T', 'C': 'G', 'G': 'C', 'T': 'A', \
                      'a': 't', 'c': 'g', 'g': 'c', 't': 'a', \
                      'u': 't', 'U': 'u', 'N': 'N', 'n': 'n'}
        complseq = [complement[base] for base in seq]
        return complseq
    rvSeq = Seq[::-1]
    return( complement_seq( rvSeq ) )

def get_seqName_in_fasta_line( line ):
    ''' 
    extract the seq name from a given fasta header line
    '''
    if( line.startswith( '>' ) ):
        # line = line.replace( '\n', '\0' )
        tag = string.split( line, ' ' )[0]
        seqN = string.strip( tag[1:] )
        return( seqN )
    else:
        return ""

def get_next_seq_Name( ifile ):
    '''
    given a fasta file, get the tag of the next seq  
    '''
    while 1:
        line = ifile.readline()
        if ( not line ):
            return False
        elif line[0] == '>':
            seqN = get_seqName_in_fasta_line( line )
            return( seqN )

def get_seqs_len( faFilename ):
    ''' get_seqs_len( faFilename) -> seqs_dict
    
    Given a fasta file, faFilename, containing many sequences,
    the function returns a dictionary, seqs_dict, with tags as keys,
    and the seq lengthes as values. 
    '''
    seqs_dict = dict()
    ifile = open( faFilename, mode = "r" )
    current_seqL = 0
    # read until the first sequence; skipping comments
    current_seqN = get_next_seq_Name( ifile )
    while 1: # collect sequence's name and 
        line = ifile.readline()
        if not line:
            if( current_seqL > 0 ):
                seqs_dict[current_seqN] = current_seqL
            break;
        elif ( line[0] == '>' ):
            if( current_seqL > 0 ):
                seqs_dict[current_seqN] = current_seqL
            current_seqN = get_seqName_in_fasta_line( line )
            current_seqL = 0;
        else:
            current_seqL = current_seqL + len( line.strip() )
    ifile.close()
    return seqs_dict

def fasta_gen( fname, seqNameSubStr ):
    ''' This function returns a generator,
    each time when the generator is called, a sequence
    with its name containing seqNameSubStr in the fasta file
    will be returned.'''
    f = open( fname )
    seq_lines = []
    title = ''

    for line in f:
        line = line.strip()
        if line.startswith( '>' ):
            if seq_lines:
                yield title, ''.join( seq_lines )
            seq_lines = []
            if( seqNameSubStr in line ):
                title = line[1:]
            else:
                title = ''
        elif title != '':
            seq_lines.append( line )
    if title != '':
        seq_lines.append( line )
        yield title, ''.join( seq_lines )
    while( 1 ):
        yield ( '', '' )
    return

def get_tag_seq_from_fasta_w_substr( faFile, seqNameSubStr ):
    '''
    Given a fasta file, return the first sequence with its
    name starts with seqNameSubStr. To get the next sequence 
    call get_tag_seq_from_fasta_w_substr( faFile, '>')
    '''
    [tag, seq] = ["", ""]
    while 1:
        line = faFile.readline()
        if not line:
            break;
        elif ( line[0] == '>' ):
            if( tag == "" ):
                if( seqNameSubStr in line ):
                    tag = string.strip( line[1:] )
            else:      # back to previous line
                #try: os.SEEK_CUR = 1
                faFile.seek( -1 * ( len( line ) + 1 ), 1 )
                #except AttributeError:
                    #print "Error with:" + line
                break; # end of the extracted sequence                
        elif ( tag != "" ):
            seq = seq + string.strip( line )
    return [tag, seq]

def print_fasta( out_fileN, seqName, seq ):
    ''' 
    extract the seq from a given fasta
    '''
    NT_PER_LINE = 80
    ofile = open( out_fileN, "a" )
    ofile.write( ">" + seqName + "\n" )
    for i in range( 0, len( seq ) ):
        ofile.write( seq[i] )
        if( i % NT_PER_LINE == NT_PER_LINE - 1 ):
            ofile.write( "\n" )
    if( i % NT_PER_LINE != 0 ):
        ofile.write( "\n" )
    ofile.close()

def get_seqs_from_fasta_w_substr\
    ( faFilename, seqNameSubStr, bSingleSeq, outFileN ):
    '''        
    Given a fasta file containing many sequences, the function 
    print the sequences with their tag containing the seqNameSubStr 
    as their substrings. 
    If bSingleSeq == True, print one sequence and exit.
    '''
    noOfSeqs = 0
    ifile = open( faFilename, mode = "r" )
    ofile = open ( outFileN, mode = "w" )
    # read until the first sequence header 
    while 1:
        line = ifile.readline()
        if ( not line ) or line == '\n' or line[0] == '>':
            break;

    # print the sequence's name and its tag if the tag contains seqNameSubStr   
    while 1:
        if not line:
            break;
        elif ( line[0] == '>' ) and ( seqNameSubStr in line ):
            noOfSeqs = noOfSeqs + 1
            if( noOfSeqs > 1 and bSingleSeq ):
                break
            ofile.write( line )
            while 1:
                line = ifile.readline()
                if ( not line ) or line == '\n' or line[0] == '>':
                    break;
                else:
                    ofile.write( line )
        else:
            line = ifile.readline()
            pass
    # end of test
    ifile.close()
    ofile.close()
    return noOfSeqs
